import cv2
import os
import numpy as np
from PIL import Image
# 显示结果（如果用Jupyter Notebook）
from matplotlib import pyplot as plt

def process_red_channel(image_path):
    # 读取图像并转换通道顺序
    pil_image = Image.open(image_path)
    img_array = np.array(pil_image)
    img = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
    # img[:, :, 2] = np.where(img[:, :, 2] < 150, 0, img[:, :, 2])
    # 调整尺寸
    new_size = (1200, 800)
    img = cv2.resize(img, new_size)
    
    # 转换到HSV颜色空间
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    
    # 定义红色范围（HSV空间）
    lower_red1 = np.array([0, 50, 50])     # 色调低区间1
    upper_red1 = np.array([10, 255, 255])  # 色调高区间1
    lower_red2 = np.array([160, 50, 50])   # 色调低区间2（环状色轮的另一端）
    upper_red2 = np.array([180, 255, 255]) # 色调高区间2
    
    # 创建红色掩膜
    mask1 = cv2.inRange(hsv, lower_red1, upper_red1)
    mask2 = cv2.inRange(hsv, lower_red2, upper_red2)
    red_mask = cv2.bitwise_or(mask1, mask2)
    
    # 形态学操作（可选，增强掩膜效果）
    kernel = np.ones((10,10), np.uint8)
    red_mask = cv2.morphologyEx(red_mask, cv2.MORPH_CLOSE, kernel)
    red_mask = cv2.morphologyEx(red_mask, cv2.MORPH_OPEN, kernel)
    
    # 应用掩膜：保留红色区域，其他区域设为黑色
    result = cv2.bitwise_and(img, img, mask=red_mask)
    result_rgb = cv2.cvtColor(result, cv2.COLOR_BGR2RGB)

    # 转换为灰度图（单通道）
    gray_result = cv2.cvtColor(result_rgb, cv2.COLOR_RGB2GRAY)
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (1, 5))
    thresh = cv2.morphologyEx(gray_result, cv2.MORPH_OPEN, kernel)
    contours, hierarchy = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    min_distance = 20
    selected_contours = []

    # 筛选轮廓并过滤小面积区域
    for cnt in contours:
        x, y, w, h = cv2.boundingRect(cnt)
        if w * h < 3000:  # 过滤面积过小的轮廓
            continue
        valid = True
        for (sx, sy, sw, sh) in selected_contours:
            if (abs(x - sx) < min_distance and abs(y - sy) < min_distance):
                valid = False
                break
        if valid:
            selected_contours.append((x, y, w, h))

    for cnt in selected_contours:
        x, y, w, h = cnt
        area = w * h
        if area > 100:
            print("xywh", x, y, w, h)
            print("矩形面积：", area)
            roi = gray_result[y:y + h, x:x + w]
            cv2.imshow("ROI", roi)
            cv2.waitKey(0)
            print("矩形面积：", area)
    # return gray_result

for file in os.listdir("images"):
    # 使用示例
    processed_image = process_red_channel(f"images/{file}")
    # plt.imshow(processed_image)
    # plt.imshow(processed_image, cmap='gray')
    # plt.axis('off')
    # plt.show()
    break